import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import os


def plotAirfoilGeom(coordinates, file_path=None):
    # plt.figure(figsize=(20, 20))

    # print(coordinates)
    len = int(coordinates.shape[0]/2)
    plt.scatter(coordinates[0:len, 0], coordinates[0:len, 1],
                color='red',
                s=10)  # , linestyle='-')
    plt.scatter(coordinates[len::, 0], coordinates[len::, 1],
                color='blue',
                s=10)  # , linestyle='-')

    # draw the arrow
    # for i in range(coordinates.shape[0] - 1):
    #     e = coordinates[i + 1]
    #     s = coordinates[i]
    #     plt.annotate("",
    #                  xy=(e[0], e[1]),
    #                  xytext=(s[0], s[1]),
    #                  arrowprops=dict(arrowstyle="->", color="r"))

    # plt.xlim((0, 1.0))  # 设置横坐标范围
    # plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
    plt.plot([0, 0.01], [0, 0], '--')
    plt.plot([0, 0], [0.001, -0.01], '--')
    plt.xlim((-0.0001, 0.1))  # 设置横坐标范围
    plt.ylim((-0.2, 0.2))  # 设置纵坐标范围
    # plt.xticks(np.arange(0.7, 1.05, 0.1))
    # plt.gca().set_aspect(1)
    # plt.grid()  # 设置网格
    # plt.show()
    plt.savefig(file_path, dpi=600, bbox_inches='tight')
    plt.close()





def csv_to_numpy(file):
    df = pd.read_csv(file)
    tmp = np.zeros(shape=df.shape)
    tmp[:, 0] = df['x'].to_numpy()
    tmp[:, 1] = df['y'].to_numpy()
    return tmp


def csv_to_numpy_v2(file):
    df = pd.read_csv(file, header=None, index_col=None)
    return df.to_numpy()


def draw_v2(seq):
    '''
        draw the arrows.
    :param seq: list of coordinate; type: numpy.ndarray; shape: [N, 2]
    :return:
    '''
    fig = plt.figure(figsize=(1, 5))
    ax = fig.add_subplot(111)

    # draw the f(x)
    x = np.arange(0, 5.05, 0.05)
    y = f(x)
    plt.plot(x, y)

    # draw the arrow
    for i in range(seq.shape[0] - 1):
        e = seq[i + 1]
        s = seq[i]
        ax.annotate("",
                    xy=(e[0], e[1]),
                    xytext=(s[0], s[1]),
                    arrowprops=dict(arrowstyle="->", color="r"))

    ax.set_xlim(0, 2)
    ax.set_ylim(-3, 3)
    ax.grid()
    ax.set_aspect('equal')  # x轴y轴等比例

    # x轴y轴等比例
    plt.show()


if __name__ == "__main__":
    file_name = r"\clarkyh.csv"

    # coordinates = csv_to_numpy(r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_1_seq_all"+file_name)
    # plotAirfoilGeom(coordinates)
    # coordinates = csv_to_numpy(r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_2_clip"+file_name)
    # plotAirfoilGeom(coordinates)
    # coordinates = csv_to_numpy(r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_3_clip"+file_name)
    # plotAirfoilGeom(coordinates)
    # coordinates = csv_to_numpy(r"D:\data_cmp\airfoil_csv_op\airfoil_predata1"+file_name)
    # plotAirfoilGeom(coordinates)
    # coordinates = csv_to_numpy(r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi"+file_name)
    # plotAirfoilGeom(coordinates)

    # dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos" # r"NACA4_proc"
    # save_dir = r"ttt"  # r"NACA4_show"
    # file = "naca63206.csv"
    # file = "bw050209.csv"
    # file_path = os.path.join(dir, file)
    # save_path = os.path.join(save_dir, file[:-3]+"jpg")
    # coordinates = csv_to_numpy(file_path)
    # plotAirfoilGeom(coordinates, save_path)

    # 数据集可视化
    # dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_7_sspv2" # r"NACA4_proc"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_7_sspv2_show"  # r"NACA4_show"
    # dir = r"D:\pycharm_workspace\gitee\CST\CST_output" # r"NACA4_proc"
    # save_dir = r"D:\pycharm_workspace\gitee\CST\CST_output_show"  # r"NACA4_show"
    # dir = r"D:\data_cmp\airfoil_csv_op\NACA4_ssp"  # r"NACA4_proc"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\NACA4_ssp_show"  # r"NACA4_show"
    dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_11_cbi"  # r"ideal_airfoil1"
    save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_11_cbi_show"  # r"ideal_airfoil1_show"
    for file in os.listdir(dir):
    # for file in ["e558.csv"]:

        file_path = os.path.join(dir, file)
        save_path = os.path.join(save_dir, file[:-3] + "jpg")
        # coordinates = csv_to_numpy_v2(file_path)
        coordinates = pd.read_csv(file_path, header=0, index_col=None).to_numpy()
        plotAirfoilGeom(coordinates, save_path)


